"""
针对不同类型的页面，速率限制的设置应该根据页面的敏感度、资源消耗和预期使用模式来调整。以下是对登录、注册和普通页面的速率限制建议：

登录页面的速率限制
登录是安全敏感操作，需要防止暴力破解，同时保证正常用户体验：

time_window: 10-30秒
max_requests: 3-5次
这样设置可以有效防止暴力破解攻击，同时允许用户在输错密码后立即重试几次。已有的登录限制器（5次/5分钟）是针对失败尝试的更长期限制，而这个是针对页面访问频率的短期限制。

注册页面的速率限制
注册通常涉及数据库写入和可能的邮件发送，且通常不需要快速重复操作：

time_window: 20-60秒
max_requests: 2-3次
相对严格的限制可以防止批量注册机器人，保护系统资源。

普通页面的速率限制
对于一般内容展示页面，可以相对宽松：

time_window: 5秒
max_requests: 10-15次
对于API或AJAX请求较多的交互页面：

time_window: 5秒
max_requests: 20-30次
资源密集型页面的速率限制
对于涉及复杂查询或大量资源加载的页面：

time_window: 5-10秒
max_requests: 5-8次
管理后台页面的速率限制
管理后台通常访问量小但权限敏感：

time_window: 5-10秒
max_requests: 10-15次
"""

from functools import wraps
import os
import json
import hashlib
import glob
import time
import threading
from datetime import datetime
from flask import request, jsonify, flash, current_app, url_for, redirect
from flask_login import current_user
from collections import defaultdict, deque
from app.common.logutil import log, log_error

def get_request_identifier():
    """
    生成请求的唯一标识符，结合多种因素以区分共享IP下的不同用户
    
    Returns:
        str: 请求的唯一标识符
    """
    # 基本标识为IP地址
    identifier = request.remote_addr
    
    # 如果用户已登录，结合用户ID
    if hasattr(current_user, 'is_authenticated') and current_user.is_authenticated:
        return f"{identifier}_user_{current_user.user_id}"
    
    return f"{identifier}"

class RateLimitFileManager:
    """速率限制文件管理器"""
    
    def __init__(self, directory="logs/rate_limits"):
        """
        初始化文件管理器
        
        Args:
            directory (str): 文件存储目录
        """
        self.directory = directory
        # 确保目录存在
        os.makedirs(directory, exist_ok=True)
    
    def get_filename(self, identifier):
        """
        根据标识符和当前时间生成文件名
        
        Args:
            identifier (str): 请求标识符
            
        Returns:
            str: 文件路径
        """
        # 哈希标识符
        hashed = hashlib.md5(identifier.encode()).hexdigest()[:10]
        # 使用当前时间戳作为文件名的一部分
        timestamp = int(time.time())
        return os.path.join(self.directory, f"rate_limit_{hashed}_{timestamp}.json")
    
    def save_limit_state(self, identifier, data):
        """
        保存限制状态到文件
        
        Args:
            identifier (str): 请求标识符
            data (dict): 状态数据
        
        Returns:
            str: 文件路径
        """
        filename = self.get_filename(identifier)
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False)
        return filename
    
    def get_limit_state(self, identifier):
        """
        获取标识符的最新限制状态
        
        Args:
            identifier (str): 请求标识符
            
        Returns:
            dict|None: 限制状态数据，不存在或已过期返回None
        """
        # 哈希标识符
        hashed = hashlib.md5(identifier.encode()).hexdigest()[:10]
        # 查找匹配文件
        pattern = os.path.join(self.directory, f"rate_limit_{hashed}_*.json")
        files = glob.glob(pattern)
        
        if not files:
            return None
            
        # 按修改时间排序，获取最新的文件
        latest_file = max(files, key=os.path.getmtime)
        
        try:
            with open(latest_file, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 检查是否过期
            if data.get("expires_at", 0) < time.time():
                # 过期了，删除文件
                os.remove(latest_file)
                return None
                
            return data
        except (json.JSONDecodeError, IOError):
            # 文件损坏或无法访问
            return None
    
    def delete_limit_files(self, identifier):
        """
        删除标识符相关的所有限制文件
        
        Args:
            identifier (str): 请求标识符
        """
        # 哈希标识符
        hashed = hashlib.md5(identifier.encode()).hexdigest()[:10]
        # 查找并删除匹配文件
        pattern = os.path.join(self.directory, f"rate_limit_{hashed}_*.json")
        for file_path in glob.glob(pattern):
            try:
                os.remove(file_path)
            except OSError:
                pass
    
    def update_verification(self, identifier, verified=True):
        """
        更新验证状态
        
        Args:
            identifier (str): 请求标识符
            verified (bool): 是否已验证
            
        Returns:
            dict: 更新后的状态数据
        """
        # 获取当前状态或创建新状态
        data = self.get_limit_state(identifier) or {}
        
        # 更新状态数据
        data.update({
            "identifier": identifier,
            "verification_status": verified,
            "verify_time": int(time.time()),
            "expires_at": int(time.time() + 300)  # 5分钟过期
        })
        
        # 删除旧文件并保存新状态
        self.delete_limit_files(identifier)
        self.save_limit_state(identifier, data)
        
        return data
    
    def clean_old_files(self, days=2):
        """
        清理指定天数前的限制文件
        
        Args:
            days (int): 清理天数，默认2天
        """
        cutoff_time = time.time() - (days * 24 * 60 * 60)
        
        # 查找所有限制文件
        pattern = os.path.join(self.directory, "rate_limit_*.json")
        for filepath in glob.glob(pattern):
            try:
                # 获取文件修改时间
                if os.path.getmtime(filepath) < cutoff_time:
                    os.remove(filepath)
            except OSError:
                continue

# 创建文件管理器实例
file_manager = RateLimitFileManager()

class IPRateLimiter:
    """IP访问速率限制器"""
    
    def __init__(self, whitelist:list=None):
        # 请求标识符 -> 访问时间队列的映射
        self.requests = defaultdict(lambda: deque(maxlen=100))
        # 默认限制：5秒内最多5次请求
        self.default_time_window = 5  # 5秒
        self.default_max_requests = 5  # 5次请求

        # 设置不同类型页面的速率限制
        self.ratedict = {
            'login': (30, 5),      # 登录页面
            'register': (60, 3),   # 注册页面
            'normal': (5, 15),     # 普通页面
            'api': (5, 30),        # API请求
            'resource': (10, 8),   # 资源密集型页面
            'admin': (10, 15)      # 管理后台页面
        }

        # 添加锁字典，为每个标识符维护一个锁
        self.locks = defaultdict(threading.RLock)
        # 全局锁，用于保护锁字典和其他共享资源
        self.global_lock = threading.RLock()
            
        # 更新配置中的速率限制设置（如果有）
        self.update_config()

        self.start_cleanup_daemon() # 启动清理后台线程

    def update_config(self):
        """从应用配置更新限制设置"""
        if current_app:
            config_limits = current_app.config.get('RATELIMIT', {})
            if config_limits:
                self.ratedict.update(config_limits)

    def start_cleanup_daemon(self):
        """启动定期清理后台线程"""
        def cleanup_task():
            while True:
                time.sleep(3600)  # 每小时执行一次
                try:
                    self.clean_expired_records()
                    # 也清理文件
                    file_manager.clean_old_files()
                except Exception as e:
                    # 确保后台线程不会因异常而退出
                    log_error(f"清理速率限制记录时出错: {str(e)}", e)
        
        cleanup_thread = threading.Thread(target=cleanup_task, daemon=True)
        cleanup_thread.start()
    
    def is_limited(self, identifier, time_window=None, max_requests=None, is_ip_only=False):
        """
        检查请求是否超过限制
        
        Args:
            identifier: 请求标识符(可以是IP或组合标识符)
            time_window: 时间窗口(秒)
            max_requests: 最大请求次数
            is_ip_only: 是否仅使用IP进行限制
            
        Returns:
            bool: 是否超过限制
        """
        # 清理过期记录
        self.clean_expired_records()

        # 获取或创建此标识符的锁
        with self.global_lock:
            lock = self.locks[identifier]

        with lock:

            # 使用默认值或传入的参数
            time_window = time_window or self.default_time_window
            max_requests = max_requests or self.default_max_requests
            
            current_time = time.time()
            request_times = self.requests[identifier]
            
            # 添加当前请求时间
            request_times.append(current_time)
            
            # 仅保留时间窗口内的请求
            while request_times and request_times[0] < current_time - time_window:
                request_times.popleft()
            
            # 检查是否超过最大请求数
            threshold = max_requests
            if is_ip_only and '.' in identifier and '_' not in identifier:
                # 共享IP环境下适当放宽限制
                threshold = max_requests * 2
            
            is_limited = len(request_times) > threshold
        
        # 如果超限，创建限制文件
        if is_limited:
            self._create_limit_file(identifier, request) # 分离IO操作
        
        return is_limited
    
    # 分离IO操作
    def _create_limit_file(self, identifier, request_obj):
        """创建限制文件（无需锁保护）"""
        data = {
            "identifier": identifier,
            "ip": identifier.split('_')[0] if '_' in identifier else identifier,
            "limited_at": int(time.time()),
            "expires_at": int(time.time() + 60),
            "verification_status": False,
            "original_url": request_obj.url,
            "captcha_code": "0",
        }
        file_manager.save_limit_state(identifier, data)
    
    def clean_expired_records(self, expiry_time=3600):
        """清理过期的请求记录"""
        current_time = time.time()
        expired_identifiers = []

        with self.global_lock:
        
            for identifier, times in self.requests.items():
                if not times or (current_time - times[-1]) > expiry_time:
                    expired_identifiers.append(identifier)
            
            for identifier in expired_identifiers:
                del self.requests[identifier]
                # 同时清理不再需要的锁
                if identifier in self.locks:
                    del self.locks[identifier]
            
    def reset(self, identifier):
        """
        重置请求标识符的计数，同时删除对应的限制文件
        
        Args:
            identifier: 请求标识符
        """
        with self.global_lock:
            # 清除内存中的记录
            if identifier in self.requests:
                del self.requests[identifier]

            # 清理此标识符的锁
            if identifier in self.locks:
                del self.locks[identifier]
            
        # 删除对应的文件
        file_manager.delete_limit_files(identifier)


# 创建全局速率限制器实例
ip_rate_limiter = IPRateLimiter()

def rate_limit(type:str="normal", time_window:int=None, max_requests:int=None, message:str="请求过于频繁，请稍后再试"):
    """
    速率限制装饰器
    
    Args:
        type: 速率限制类型, 可选值: login, register, normal, api, resource, admin
        time_window: 时间窗口(秒), 如果为None则根据类型获取
        max_requests: 时间窗口内允许的最大请求数, 如果为None则根据类型获取
        message: 超出限制时的提示信息
        
    Returns:
        装饰器函数
    """
    def decorator(f):
        @wraps(f)
        def wrapped_function(*args, **kwargs):
            # 确保配置是最新的
            ip_rate_limiter.update_config()

            # 获取客户端IP和请求标识符
            client_ip = request.remote_addr
            request_id = get_request_identifier()
            
            # 是否在IP白名单中，如果在白名单中则不限制
            _whitelist = current_app.config.get('IP_WHITELIST', [])
            if client_ip in _whitelist:
                return f(*args, **kwargs)
            
            # 确定用户信息，用于日志记录
            _user = "未登录"
            if hasattr(current_user, 'is_authenticated') and current_user.is_authenticated:
                _user = f"{current_user.loginname}-({current_user.user_id}::{current_user.nickname})"
            
            # 检查是否已存在限制状态且未验证
            limit_state = file_manager.get_limit_state(request_id) 
            if limit_state: # 文件存在就说明已被限制
                # 已被限制且未验证，重定向到验证页面
                _mess = f"{client_ip} ({_user}) 请求过于频繁，请进行人机验证。"
                log_error(f"{client_ip} ({_user}) 请求过于频繁，已被限制", level='warning')

                flash(_mess, 'error')

                return redirect(url_for('user.verify_human', next=request.url))
            
            # 根据类型获取速率限制参数
            _time_window, _max_requests = time_window, max_requests
            if not _time_window or not _max_requests:
                _time_window, _max_requests = ip_rate_limiter.ratedict.get(type, (5, 15))
                if not _time_window:
                    _time_window = time_window or 5
                if not _max_requests:
                    _max_requests = max_requests or 15
            
            # 检查是否超过限制
            if ip_rate_limiter.is_limited(request_id, _time_window, _max_requests):
                _mess = f"{client_ip} ({_user}) 请求过于频繁，请稍后再试或进行人机验证。{message}"
                log_error(f"{client_ip} ({_user}) 请求过于频繁，已被限制", level='warning')
                
                # API请求返回JSON响应
                if request.is_json or request.path.startswith('/api/'):
                    return jsonify({
                        'success': False,
                        'message': _mess,
                        'code': 429  # Too Many Requests
                    }), 429
                
                # 网页请求重定向到验证页面
                flash(_mess, 'error')

                return redirect(url_for('user.verify_human', next=request.url))
            
            # 未超限，执行原始视图函数
            return f(*args, **kwargs)
        
        return wrapped_function
    
    return decorator
